import csv
from pathlib import Path


def split_csv_by_field(input_file, field_name, output_dir=None):
    """按指定字段拆分CSV文件到指定目录"""
    input_path = Path(input_file)

    # 设置输出目录（默认为输入文件同级目录下的拆分文件夹）
    if output_dir is None:
        output_dir = input_path.parent / f"{input_path.stem}_split"
    else:
        output_dir = Path(output_dir)

    # 创建输出目录（如果不存在）
    output_dir.mkdir(parents=True, exist_ok=True)

    with open(input_file, "r", encoding="utf-8") as f:
        reader = csv.DictReader(f)

        # 存储已打开的文件句柄
        file_handles = {}

        for row in reader:
            field_value = row[field_name]

            # 如果该字段值对应的文件尚未打开，则创建新文件
            if field_value not in file_handles:
                # 创建安全的文件名（替换可能无效的字符）
                safe_name = "".join(
                    c if c.isalnum() or c in (" ", "-", "_") else "_"
                    for c in str(field_value)
                ).strip()[:100]  # 限制文件名长度

                # 生成唯一文件名（避免重复）
                output_file = output_dir / f"{safe_name}.csv"
                counter = 1
                while output_file.exists():
                    output_file = output_dir / f"{safe_name}_{counter}.csv"
                    counter += 1

                # 以追加模式打开文件并写入表头
                fh = open(output_file, "a", encoding="utf-8", newline="")
                writer = csv.DictWriter(fh, fieldnames=reader.fieldnames)
                if fh.tell() == 0:  # 如果是新文件，写入表头
                    writer.writeheader()
                file_handles[field_value] = (fh, writer)

            # 写入当前行
            file_handles[field_value][1].writerow(row)

    # 关闭所有文件句柄
    for fh, _ in file_handles.values():
        fh.close()

    print(f"拆分完成，共生成 {len(file_handles)} 个文件，保存在: {output_dir}")


if __name__ == "__main__":
    # 使用示例
    input_csv = "merged_output.csv"  # 替换为你的大CSV文件路径
    split_field = "点位概况描述"  # 替换为你要拆分的字段名
    save_dir = "split_files"  # 替换为你想要保存的目录

    split_csv_by_field(input_csv, split_field, save_dir)
